from IPython.display import Markdown
import asyncio
import os
import uuid
from langgraph.checkpoint.memory import MemorySaver
from graph import builder
from langchain_community.chat_models import ChatOllama
# Set environment variables for API keys and configuration
# os.environ["TAVILY_API_KEY"] = ""

# Define the report structure
REPORT_STRUCTURE = """使用此结构来创建关于用户提供主题的报告：

1. 引言（无需研究）
   - 主题领域的简要概述

2. 主体部分：
   - 每个部分应关注用户提供主题的一个子主题
   
3. 结论
   - 目标是提供1个结构化元素（列表或表格）来提炼主体部分的要点
   - 对报告进行简明扼要的总结"""

# Define model constants
OLLAMA_MODEL_NAME = "qwen2.5:14b"

def get_ollama_model():
    """Initialize and return the Ollama model with qwen2.5:14b."""
    return ChatOllama(model=OLLAMA_MODEL_NAME)

def run_deep_research(topic: str):
    """Run deep research with local Ollama model.
    
    Args:
        topic: The research topic to investigate
    """
    # Initialize memory saver
    memory = MemorySaver()
    
    # Build the graph
    graph = builder.compile(checkpointer=memory)
    
    # Configure the thread with Ollama model
    thread = {
        "thread_id": str(uuid.uuid4()),
        "configurable": {
            "search_api": "tavily",
            "planner_model": OLLAMA_MODEL_NAME,
            "writer_model": OLLAMA_MODEL_NAME,
            "max_search_depth": 2,
            "report_structure": REPORT_STRUCTURE
        }
    }

    # Set up initial state
    initial_state = {
        "topic": topic,
        "config_params": thread["configurable"]
    }
    
    # Run the graph and get the result
    result = asyncio.run(graph.ainvoke(initial_state, thread))
    
    # Output the result
    print("Graph execution result:")
    print(Markdown(str(result)))
    
    # Check if the result contains a finish value
    if "finish" in result:
        print("\nGenerated Report:")
        print(Markdown(result["finish"]))
    # If the result contains a report, display it nicely
    elif "report" in result:
        print("\nGenerated Report:")
        print(Markdown(result["report"]))

def main():
    """Command line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(description="DeepResearch - AI Research Assistant")
    parser.add_argument("topic", help="The research topic to investigate")
    
    args = parser.parse_args()
    run_deep_research(args.topic)

if __name__ == "__main__":
    # Example usage
    main()
